﻿/*
Create index file from data file "*-pages-articles.xml.bz2"
*/
#define DEBUG 0

#pragma comment(lib, "../bzip2/libbz2.lib")
#include "../bzip2/bzlib.h"
#include "../sqlite/sqlite3.h"
#include <stdio.h>
#include <time.h>

int UTF8toUCS2(wchar_t* dst, int count, const unsigned char* str)
{
    int i = 0, length = 0;

    if (!str || !dst) return 0;
    while (str[i] && length < count)
    {
        if (str[i] <= 0x7FU)      //ASCII
        {
            dst[length] = (wchar_t)str[i]; 
            i++;    length++; 
        }
        else if (str[i] <= 0xC1U) //part of multi-byte or overlong encoding ->ignore
        {
            i++;          
        }
        else if (str[i] <= 0xDFU) //2-byte
        {
            dst[length] = ((str[i]&0x001FU)<<6) | (str[i+1]&0x003FU); 
            i += 2; length++; 
        }
        else if (str[i] <= 0xEFU) //3-byte
        {
            dst[length] = ((str[i]&0x001FU)<<12) | ((str[i+1]&0x003FU)<<6) | (str[i+2]&0x003FU); 
            i += 3; length++; 
        }
        else
        {
            i++;                   //4-byte, restricted or invalid range ->ignore
        }
    }
    dst[length] = 0;
    return length;
}

int getText(BZFILE** pbfp, char* word, wchar_t** text)
{
    int bzerror;
    char* buf;
    char ch;
    int count = 0;
    int size = 0;
    int i = 0;

    while(word[i])
    {
        do
        {
            BZ2_bzRead(&bzerror, *pbfp, &ch, 1);
        } while(*word != ch && bzerror == BZ_OK);
        if (bzerror == BZ_STREAM_END)
        {
            return bzerror;
        }
        if (*word == ch)
        {
            while(word[++i] && bzerror == BZ_OK)
            {
                BZ2_bzRead(&bzerror, *pbfp, &ch, 1);
                if (word[i] != ch)
                {
                    goto BREAK2;
                }
            }
            break;
        }
BREAK2:
        i = 0;
    }
    buf = (char*)malloc(1024);
    if (buf == NULL)
    {
        return -10;
    }
    do
    {
        BZ2_bzRead(&bzerror, *pbfp, &ch, 1);
        buf[size++] = ch;
        if (++count >= 1024)
        {
            count = 0;
            buf = (char*)realloc(buf, size + 1024);
            if (buf == NULL)
            {
                return -10;
            }
        }
    } while(ch != '<' && bzerror == BZ_OK);
    buf[--size] = 0;
    size *= 2 + 2;
    *text = (wchar_t*)malloc(size);
    if (*text == NULL)
    {
        free(buf);
        return -10;
    }
    UTF8toUCS2(*text, size / 2, buf);
    free(buf);
    return bzerror;
}

static sqlite3 *db_title, *db_text;
static sqlite3_stmt *stmt_title = NULL, *stmt_text = NULL;
static char* errmsg;
int dbOpen(int db)
{
    time_t timer;
    struct tm *t_st;
    wchar_t db_file[64];

    time(&timer);
    t_st = localtime(&timer);
    printf("%02d:%02d:%02d DB-%d ファイルを作成中…\n", t_st->tm_hour, t_st->tm_min, t_st->tm_sec, db);
    swprintf(db_file, 64, L"text%d.db", db);
    sqlite3_open16(db_file, &db_text);
    sqlite3_exec(db_text, "CREATE TABLE wiki_text (id INTEGER, p_text BLOB, text_size INTEGER)", NULL, NULL, &errmsg);
    if (errmsg)
    {
        printf("%s\n", errmsg);
        sqlite3_free(errmsg);
    }
    sqlite3_exec(db_text, "BEGIN", NULL, NULL, &errmsg);
    if (errmsg)
    {
        printf("%s\n", errmsg);
        sqlite3_free(errmsg);
    }
    sqlite3_prepare_v2(db_text, "INSERT INTO wiki_text (id, p_text, text_size) values (?, ?, ?)", -1, &stmt_text, NULL);
    return 0;
}

int dbClose(int db)
{
    time_t timer;
    struct tm *t_st;

    sqlite3_finalize(stmt_text);
    time(&timer);
    t_st = localtime(&timer);
    printf("\n%02d:%02d:%02d DB-%d IDインデックス作成中…\n", t_st->tm_hour, t_st->tm_min, t_st->tm_sec, db);
    sqlite3_exec(db_text, "CREATE INDEX idx_id_text ON wiki_text (id)", NULL, NULL, &errmsg);
    if (errmsg)
    {
        printf("%s\n", errmsg);
        sqlite3_free(errmsg);
    }
    sqlite3_exec(db_text, "COMMIT", NULL, NULL, &errmsg);
    if (errmsg)
    {
        printf("%s\n", errmsg);
        sqlite3_free(errmsg);
    }
    sqlite3_close(db_text);
    return 0;
}

int main(int argc, char* argv[])
{
    int bzerror, db_id = 0;
    unsigned int count;
    FILE *fp;
    BZFILE* bfp;
    wchar_t *title, *id_text, *text;
    char* bzbuf;
    time_t timer;
    struct tm *t_st;

#if DEBUG
    fp = fopen("../jawiki-latest-pages-articles.xml.bz2", "rb");
#else
    if (argc < 2)
    {
        printf("Usage: >indexer.exe *-pages-articles.xml.bz2\n");
        return 0;
    }
    fp = fopen(argv[1], "rb");
#endif
    if (fp == NULL)
    {
        printf("fopen error\n");
        return 1;
    }
    bfp = BZ2_bzReadOpen(&bzerror, fp, 0, 0, NULL, 0);
    if (bzerror != BZ_OK)
    {
        printf("bzopen error %d\n", bzerror);
        BZ2_bzReadClose(&bzerror, bfp);
        fclose(fp);
        return 1;
    }
    sqlite3_open16(L"title.db", &db_title);
    sqlite3_exec(db_title, "CREATE TABLE wiki_title (id INTEGER, db_number INTEGER, p_title TEXT COLLATE NOCASE)", NULL, NULL, &errmsg);
    if (errmsg)
    {
        printf("%s\n", errmsg);
        sqlite3_free(errmsg);
        goto END;
    }
    sqlite3_exec(db_title, "BEGIN", NULL, NULL, &errmsg);
    if (errmsg)
    {
        printf("%s\n", errmsg);
        sqlite3_free(errmsg);
    }
    sqlite3_prepare_v2(db_title, "INSERT INTO wiki_title (id, db_number,  p_title) values (?, ?, ?)", -1, &stmt_title, NULL);
    dbOpen(db_id);
    count = 1;
    do
    {
        long id;
        int ret;
        size_t size;
        unsigned int bzlen;

        title = NULL;
        id_text = NULL;
        text = NULL;
        bzbuf = NULL;
        bzerror = getText(&bfp, "<title>", &title);
        if (bzerror != BZ_OK)
        {
            break;
        }
        bzerror = getText(&bfp, "<id>", &id_text);
        if (bzerror != BZ_OK)
        {
            break;
        }
        id = _wtol(id_text);
        bzerror = getText(&bfp, "<text xml:space=\"preserve\">", &text);
        if (bzerror != BZ_OK)
        {
            break;
        }
        size = wcslen(text) * 2;
        bzlen = size + 100; // for small size(<100)
        bzbuf = (char*)malloc(bzlen);
        if (bzbuf == NULL)
        {
            break;
        }
        ret = BZ2_bzBuffToBuffCompress(bzbuf, &bzlen, (char*)text, size, 9, 0, 0);
        if (ret != BZ_OK)
        {
            break;
        }
        sqlite3_bind_int(stmt_title, 1, id);
        sqlite3_bind_int(stmt_title, 2, db_id);
        sqlite3_bind_text16(stmt_title, 3, title, -1, SQLITE_TRANSIENT);
        ret = sqlite3_step(stmt_title);
        if (ret != SQLITE_DONE)
        {
            printf("step error: %s\n", sqlite3_errmsg(db_title));
            break;
        }
        ret = sqlite3_reset(stmt_title);
        if (ret != SQLITE_OK)
        {
            printf("reset error: %s\n", sqlite3_errmsg(db_title));
            break;
        }
        ret = sqlite3_clear_bindings(stmt_title);
        if (ret != SQLITE_OK)
        {
            printf("clear error: %s\n", sqlite3_errmsg(db_title));
            break;
        }
        sqlite3_bind_int(stmt_text, 1, id);
        sqlite3_bind_blob(stmt_text, 2, bzbuf, bzlen, SQLITE_TRANSIENT);
        sqlite3_bind_int(stmt_text, 3, size);
        ret = sqlite3_step(stmt_text);
        if (ret != SQLITE_DONE)
        {
            printf("step error: %s\n", sqlite3_errmsg(db_text));
            break;
        }
        ret = sqlite3_reset(stmt_text);
        if (ret != SQLITE_OK)
        {
            printf("reset error: %s\n", sqlite3_errmsg(db_text));
            break;
        }
        ret = sqlite3_clear_bindings(stmt_text);
        if (ret != SQLITE_OK)
        {
            printf("clear error: %s\n", sqlite3_errmsg(db_text));
            break;
        }
        free(title);
        free(id_text);
        free(text);
        free(bzbuf);
        title = NULL;
        id_text = NULL;
        text = NULL;
        bzbuf = NULL;
        printf("\r%lu", count);
        if ((count % 200000) == 0)
        {
            dbClose(db_id);
            db_id++;
            dbOpen(db_id);
        }
        count++;
    } while (bzerror == BZ_OK);
    dbClose(db_id);
    free(title);
    free(id_text);
    free(text);
    free(bzbuf);
    sqlite3_finalize(stmt_title);
    time(&timer);
    t_st = localtime(&timer);
    printf("\n%02d:%02d:%02d タイトルインデックス作成中…\n", t_st->tm_hour, t_st->tm_min, t_st->tm_sec);
    sqlite3_exec(db_title, "CREATE INDEX idx_title_title ON wiki_title (p_title)", NULL, NULL, &errmsg);
    if (errmsg)
    {
        printf("%s\n", errmsg);
        sqlite3_free(errmsg);
    }
    sqlite3_exec(db_title, "COMMIT", NULL, NULL, &errmsg);
    if (errmsg)
    {
        printf("%s\n", errmsg);
        sqlite3_free(errmsg);
    }
    if (bzerror == BZ_STREAM_END)
    {
        printf("OK\n");
    }
END:
    sqlite3_close(db_title);
    BZ2_bzReadClose(&bzerror, bfp);
    fclose(fp);
    time(&timer);
    t_st = localtime(&timer);
    printf("%02d:%02d:%02d END\n", t_st->tm_hour, t_st->tm_min, t_st->tm_sec);
    return 0;
}
